#include "../include/TrieNode.h"

#include <iostream>
using namespace std;
//实现键值映射的前缀树
class MapSum
{
private:
    TrieNode *trieRoot;
    //进行一个词后缀的添加，如果已经有这个后缀则不用这个函数
    void addWord(string &word, int val)
    {
        TrieNode *ptr = trieRoot;
        for (auto ch : word)
        {
            if (ptr->children[ch - 'a'] == NULL)
            {
                ptr->children[ch - 'a'] = new TrieNode();
            }
            ptr = ptr->children[ch - 'a'];
        }
        ptr->isKey = true;
        ptr->val = val;
    }

    //进行一个累加，递归
    int sumVal(TrieNode *ptr)
    {
        int tempSum = 0;
        if (ptr == NULL)
            return 0;
        if (ptr->isKey)
            tempSum += ptr->val;
        for (int i = 0; i < 26; i++)
        {
            if (ptr->children[i] != NULL)
                tempSum += sumVal(ptr->children[i]);
        }
        return tempSum;
    }

public:
    MapSum()
    {
        trieRoot = new TrieNode();
    }

    void insert(string key, int val)
    {
        addWord(key, val);
    }

    int sum(string prefix)
    {
        TrieNode *ptr = trieRoot; //扫描这棵树，确定prefix尾端字符所在的节点
        for (auto ch : prefix)
        {
            if (ptr->children[ch - 'a'] != NULL)
            {
                ptr = ptr->children[ch - 'a'];
            }
            else
            {
                //如果prefix中出现无法匹配的字符，则说明树中不存在此前缀
                return 0;
            }
        }
        //现在ptr在prefix尾端字符所在的节点，我们只要统计其中所有key对应的val总和即可
        return sumVal(ptr);
    }
};